/*
 * Reranker.cpp
 *
 *  Created on: Dec 24, 2014
 *      Author: Phong
 */

#include "Reranker.h"
#include "Utils.h"
#include <cmath>

using namespace std;

int* Reranker::countMatches(Depstruct* testDs, Depstruct* goldDs, bool punc) {
    int uCount = 0;
    int lCount = 0;
    int n = 0;
    
    if (testDs->nWords != goldDs->nWords) {
        Utils::error("coutMatches: size not match");
    }
    
    for (int i = 1; i < testDs->nWords; i++) {
        bool count = true;
        if (punc == false && puncList.find(goldDs->tokens[i]) != puncList.end())
            count = false;
        if (count) {
            n++;
            if (testDs->head[i] == goldDs->head[i]) {
                uCount++;
                if (testDs->deprel[i] == goldDs->deprel[i])
                    lCount++;
            }
        }
    }
    
    return new int[3] { uCount, lCount, n };
}

real* Reranker::computeScore(vector<Depstruct*> testDs, vector<Depstruct*> goldDs, bool punc) {
    int uCount = 0;
    int lCount = 0;
    int n = 0;
    
    if (testDs.size() != goldDs.size()) {
        cout << testDs.size() << " " << goldDs.size() << endl << flush;
        Utils::error("computeScore: not match");
    }
    
    for (int i = 0; i < goldDs.size(); i++) {
        int* ret = countMatches(testDs[i], goldDs[i], punc);
        uCount += ret[0];
        lCount += ret[1];
        n += ret[2];
        delete[] ret;
    }
    
    return new real[2] { (real)uCount / n , (real)lCount / n };
}

void Reranker::outputFile(vector<Depstruct*> dsbank, string path, Dictionary *posDic, Dictionary *deprelDic){
    ofstream f(path);
    for (Depstruct* ds : dsbank) {
        f << ds->toString(posDic, deprelDic);
        f << endl;
    }
    f.close();
}

vector< vector<real> > Reranker::computeLogprobs(IORNN *net, KBestTreebank *kbestTb) {
    vector< vector<real> > ret;
    ret.resize(kbestTb->trees.size());
    int step = ceil(ret.size() / (float) nThreads);
    int count = 0;
    
#pragma omp parallel for
    for (int th = 0; th < nThreads; th++)
        for (int i = th*step; i < min((th+1)*step, (int)ret.size()); i++) {
            vector<Deptree*> trees = kbestTb->trees[i];
            vector<real> lprobs = net->computeLogProb(trees);
            ret[i] = lprobs;
        
            count++;
            if (count % 100 == 0) cout << count << '\r' << flush;
        }
	cout << count << endl;
    return ret;
}

vector<Depstruct*> Reranker::rerank(KBestTreebank* kbestTb, Treebank *goldTb, string type) {
    if (kbestTb->trees.size() != goldTb->size()) {
        cout << kbestTb->trees.size() << " " << goldTb->size() << flush;
        Utils::error("rerank: not match");
    }
    
    vector<Depstruct*> best(goldTb->size());
    
    for (int i = 0; i < goldTb->size(); i++) {
        Depstruct* bestDs = NULL;
        int bestU = -1;
        int bestL = -1;
        vector<Deptree*> candidates = kbestTb->trees[i];
        
        for (int j = 0; j < candidates.size(); j++) {
            int* ret = countMatches(candidates[j]->ds, goldTb->storage[i]->ds, countPunc);
            int u = ret[0];
            int l = ret[1];
            delete[] ret;
            
            if (type == "best" && (bestDs == NULL || u > bestU)) {
                bestU = u;
                bestL = l;
                bestDs = candidates[j]->ds->share();
            }
            
            else if (type == "worst" && (bestDs == NULL || u < bestU)) {
                bestU = u;
                bestL = l;
                bestDs = candidates[j]->ds;
            }
            
            else if (type == "first") {
                bestU = u;
                bestL = l;
                bestDs = candidates[j]->ds;
                break;
            }
        }
        
        best[i] = bestDs;
    }
    
    return best;
}

vector<Depstruct*> Reranker::rerank(vector< vector<real> > iornnscores, vector< vector<real> > pscores, int kbest, real alpha,
                                 KBestTreebank* kbestTb) {
    vector<Depstruct*> best(kbestTb->trees.size());
    
    for (int i = 0; i < kbestTb->trees.size(); i++) {
        Depstruct* bestDs = NULL;
        real bestScore = -1e10;
        vector<Deptree*> candidates = kbestTb->trees[i];
        vector<real> score1 = iornnscores[i];
        vector<real> score2 = pscores[i];
        
        for (int j = 0; j < min((int)candidates.size(), kbest); j++) {
            real score = score1[j] * alpha + score2[j] * (1-alpha);
            if (bestDs == NULL || score > bestScore) {
                bestDs = candidates[j]->ds;
                bestScore = score;
            }
        }
        
        best[i] = bestDs;
    }
    
    return best;
}

real* Reranker::optimizeMixReranker(IORNN *net, Treebank *devTb, KBestTreebank *kbestDevTb) {
    // compute iornn scores
    cout << "compute log probs" << endl;
    vector< vector<real> > iornnscores = computeLogprobs(net, kbestDevTb);
    vector<Depstruct*> goldDs = devTb->getDsBank();
    vector< vector<real> > pscores = kbestDevTb->scores;
    real* bestScore = NULL;
    
    cout << "k\talpha\tUAS\tLAS" << endl;
    
    for (int k = 1; k <= nTop; k++) {
        real bestAlpha = -1;
        real bestUAS = -1;
        real bestLAS = -1;
        for (real alpha = 0; alpha <= 1.; alpha += 0.005) {
            vector<Depstruct*> best = rerank(iornnscores, pscores, k, alpha, kbestDevTb);
            real* as = computeScore(best, goldDs, countPunc);
            if (bestUAS < as[0]) {
                bestAlpha = alpha;
                bestUAS = as[0];
                bestLAS = as[1];
            }
            delete[] as;
        }
        printf("%d\t%f\t%.2f\t%.2f\n", k, bestAlpha, bestUAS*100, bestLAS*100);
        
        if (bestScore == NULL)
            bestScore = new real[2] {bestUAS, bestLAS};
        else if (bestScore[0] < bestUAS) {
            bestScore[0] = bestUAS;
            bestScore[1] = bestLAS;
        }
    }
    return bestScore;
}

void Reranker::parse(IORNN* net, Treebank* testTb, KBestTreebank* kbestTestTb, int kbest, real alpha, string output) {
    vector< vector<real> > iornnscores = computeLogprobs(net, kbestTestTb);
    vector<Depstruct*> goldDs = testTb->getDsBank();
    vector< vector<real> > pscores = kbestTestTb->scores;
    vector<Depstruct*> best = rerank(iornnscores, pscores, kbest, alpha, kbestTestTb);
    real* as = computeScore(best, goldDs, countPunc);
    cout << "UAS = " << as[0] << " \t LAS = " << as[1] << endl;
    delete[] as;
    
    ofstream f(output);
    for (Depstruct* ds : best) {
        f << ds->toString(net->posDic, net->deprelDic) << endl;
    }
    f.close();
}




